# import NpyIter as NI
from keras.models import Sequential
from keras.models import Graph
from keras.layers.core import Activation, Dense, Merge, Permute, Dropout
import numpy as np



def predict(state):
    model = Graph()
    from keras.layers.core import Dense, Activation
    all=np.load('model/weight_340.h5.npy')

    model.add_input(name='input',input_shape = (361,))

    model.add_node(Dense(output_dim=361,input_dim=361),name='B1',input='input')
    model.add_node(Activation('tanh'),name='B1_tanh',input='B1')
    model.add_node(Dense(output_dim=361,input_dim=361),name='B2',input='input')
    model.add_node(Activation('tanh'),name='B2_tanh',input='B2')
    model.add_node(Dense(output_dim=361,input_dim=361),name='B', inputs=['B1_tanh', 'B2_tanh'], merge_mode='mul')
    # model.add_node(Activation('tanh'),name='B_tanh',input='B')
    model.add_node(Dense(output_dim=500,input_dim=361),name='C1_B',input='B')
    model.add_node(Dense(output_dim=500,input_dim=361),name='C1_x',input='input')
    model.add_node(Dense(output_dim=500,input_dim=500),name='C1', inputs=['C1_B', 'C1_x'], merge_mode='sum')
    model.add_node(Activation('tanh'),name='C1_tanh',input='C1')
    model.add_node(Dense(output_dim=500,input_dim=361),name='C2_B',input='B')
    model.add_node(Dense(output_dim=500,input_dim=361),name='C2_x',input='input')
    model.add_node(Dense(output_dim=500,input_dim=500),name='C2', inputs=['C2_B', 'C2_x'], merge_mode='sum')
    model.add_node(Activation('tanh'),name='C2_tanh',input='C2')
    model.add_node(Dense(output_dim=500,input_dim=500),name='C', inputs=['C1_tanh', 'C2_tanh'], merge_mode='mul')


    model.add_node(Dense(output_dim=1000,input_dim=500),name='D1',input='C')
    model.add_node(Activation('tanh'),name='D1_tanh',input='D1')
    model.add_node(Dense(output_dim=1000,input_dim=500),name='D2',input='C')
    model.add_node(Activation('tanh'),name='D2_tanh',input='D2')
    model.add_node(Dense(output_dim=1000,input_dim=1000),name='D', inputs=['D1_tanh', 'D2_tanh'], merge_mode='mul')


    model.add_node(Dense(output_dim=500,input_dim=1000),name='E1_B',input='D')
    model.add_node(Dense(output_dim=500,input_dim=500),name='E1_x',input='C')
    model.add_node(Dense(output_dim=500,input_dim=500),name='E1', inputs=['E1_B', 'E1_x'], merge_mode='sum')
    model.add_node(Activation('tanh'),name='E1_tanh',input='E1')
    model.add_node(Dense(output_dim=500,input_dim=1000),name='E2_B',input='D')
    model.add_node(Dense(output_dim=500,input_dim=500),name='E2_x',input='C')
    model.add_node(Dense(output_dim=500,input_dim=500),name='E2', inputs=['E2_B', 'E2_x'], merge_mode='sum')
    model.add_node(Activation('tanh'),name='E2_tanh',input='E2')
    model.add_node(Dense(output_dim=361,input_dim=500),name='E', inputs=['E1_tanh', 'E2_tanh'], merge_mode='mul')

    model.add_output(name='output',input='E')
    from keras.optimizers import SGD
    # model.compile(optimizer=SGD(lr=0.01, momentum=0.9, nesterov=True),{'output':'mse'})
    model.compile('rmsprop', {'output':'mse'})
    j=0
    for i in range(0,len(model.layers)):
    	if model.layers[i].get_weights()!=[]:
    		if j>38:
    			# print 'OOPS....ERROR'
    			break
    		model.layers[i].set_weights([all[j],all[j+1]])
    		j=j+2


    # test_dataiter = NI.NpyIter(
    #     	root_dir             = "./",
    #     	flist_name           = "demo.lst",
    #         batch_size = 1
    # 	)
    # print "this is in predict"
    predictions = model.predict({'input':state.reshape(1,361)})
    b = predictions.get('output')
    _positon = []
    for i in range(361):
        index = np.argmax(abs(b))
        _positon.append(index)
        b[index] = 0
    return _positon
